#include <opencv2/opencv.hpp>
#include <dlib/opencv.h>
#include <dlib/image_processing/frontal_face_detector.h>
#include <dlib/image_processing.h>
#include <dlib/image_io.h>
#include <nadjieb/mjpeg_streamer.hpp>
#include <unordered_map>
#include <unordered_set>
#include <map>
#include <sys/stat.h>
#include <chrono>
#include <iomanip>
#include <ctime>
#include <csignal>
#include <thread>
#include <mutex>
#include <queue>
#include <atomic>
#include <condition_variable>
#include <fstream>
#include <filesystem>
#include "ConfigParser.h"
#include "FaceRecognizer.h"
#include "PerformanceTimer.h"

struct RecognitionTask {
    int frame_id;
    dlib::matrix<dlib::rgb_pixel> image;
};

struct RecognitionResult {
    int frame_id;
    std::vector<dlib::rectangle> faces;
    std::vector<std::string> identities;
};

std::queue<RecognitionTask> task_queue;
std::mutex task_mutex;
std::condition_variable task_cv;

std::queue<RecognitionResult> result_queue;
std::mutex result_mutex;
std::condition_variable result_cv;

std::atomic<bool> stop_threads(false);

struct VisitorInfo {
    std::chrono::system_clock::time_point first_seen;
    std::chrono::system_clock::time_point last_seen;
    std::chrono::system_clock::time_point departure_time;
    bool present = true;
    std::string img_path;
};

std::mutex visitor_mutex;
std::unordered_map<std::string, VisitorInfo> visitor_log;

void recognition_worker(FaceRecognizer* recognizer) {
    TimerManager& timer_manager = TimerManager::getInstance();
    while (true) {
        RecognitionTask task;
        {
            std::unique_lock<std::mutex> lock(task_mutex);
            task_cv.wait(lock, []{ return !task_queue.empty() || stop_threads.load(); });
            if (stop_threads.load()) break;
            if (task_queue.empty()) continue;
            task = std::move(task_queue.front());
            task_queue.pop();
        }
        timer_manager.start("人脸检测");
        auto faces = recognizer->detectFaces(task.image);
        timer_manager.end("人脸检测");

        std::vector<std::string> identities;
        for (const auto& face : faces) {
            try {
                timer_manager.start("人脸校正");
                auto face_chip = recognizer->alignFace(task.image, face);
                timer_manager.end("人脸校正");
                timer_manager.start("人脸匹配");
                std::string identity = recognizer->recognizeFace(face_chip);
                timer_manager.end("人脸匹配");
                identities.push_back(identity);
            } catch (const std::exception& e) {
                std::cerr << "人脸处理错误: " << e.what() << std::endl;
                identities.push_back("Error");
            }
        }
        RecognitionResult result{task.frame_id, std::move(faces), std::move(identities)};
        {
            std::lock_guard<std::mutex> lock(result_mutex);
            result_queue.push(std::move(result));
        }
        result_cv.notify_one();
    }
}

volatile sig_atomic_t signal_received = 0;
void signal_handler(int signal) { signal_received = signal; }

int main() {
    std::signal(SIGINT, signal_handler);
    std::signal(SIGTERM, signal_handler);

    std::filesystem::create_directories("visitor_images");

    TimerManager& timer_manager = TimerManager::getInstance();
    ConfigParser config;
    if (!config.load("/home/ubuntu/dlib/myapp/config.txt")) {
        std::cerr << "配置文件加载失败" << std::endl;
        return -1;
    }
    if (config.get("debug") == "true") config.printAll();
    int sample_interval = 5;
    std::string interval_str = config.get("sample_interval");
    if (!interval_str.empty()) 
    {
        try { sample_interval = std::stoi(interval_str); }
        catch (...) { std::cerr << "无效的sample_interval配置，使用默认值5" << std::endl; }
    }
    std::cout << "人脸识别间隔: " << sample_interval << " 帧" << std::endl;
    FaceRecognizer face_recognizer(config);
    std::cout << "人脸识别器初始化完成" << std::endl;
    face_recognizer.printFaceLibInfo();

    const int num_workers = 2;
    std::vector<std::thread> workers;
    for (int i = 0; i < num_workers; i++) 
    {
        workers.emplace_back(recognition_worker, &face_recognizer);
        std::cout << "启动人脸识别工作线程 " << i + 1 << std::endl;
    }

    cv::VideoCapture cap;
    bool use_camera = (config.get("use_camera") == "true");
    bool video_loop = true;
    if (use_camera) 
    {
        cap.open(0);
        if (!cap.isOpened()) 
        {
            std::cerr << "无法打开摄像头" << std::endl;
            return -1;
        }
    } 
    else 
    {
        std::string source_path = config.get("video_path");
        if (source_path.empty()) 
        {
            std::cerr << "错误: 未指定视频路径 (video_path)" << std::endl;
            return -1;
        }
        cap.open(source_path);
        if (!cap.isOpened()) 
        {
            std::cerr << "无法打开视频文件: " << source_path << std::endl;
            return -1;
        }
    }

    nadjieb::MJPEGStreamer streamer;
    streamer.start(8080);

    std::vector<int> params = {cv::IMWRITE_JPEG_QUALITY, 90};
    int frame_count = 0;
    int last_processed_id = -1;
    std::vector<dlib::rectangle> last_faces;
    std::vector<std::string> last_identities;

    while (!signal_received) 
    {
        timer_manager.start("帧捕获和处理全过程");
        cv::Mat frame;
        timer_manager.start("帧捕获");
        cap >> frame;
        timer_manager.end("帧捕获");
        if (frame.empty()) 
        {
            if (use_camera) break;
            else if (video_loop) 
            {
                cap.set(cv::CAP_PROP_POS_FRAMES, 0);
                cap >> frame;
                if (frame.empty()) break;
            } 
            else 
                break;
        }

        dlib::matrix<dlib::rgb_pixel> dlib_img;
        dlib::assign_image(dlib_img, dlib::cv_image<dlib::bgr_pixel>(frame));

        RecognitionResult result;
        bool new_result = false;
        {
            std::unique_lock<std::mutex> lock(result_mutex);
            if (result_cv.wait_for(lock, std::chrono::milliseconds(1), []{ return !result_queue.empty(); })) 
            {
                result = std::move(result_queue.front());
                result_queue.pop();
                new_result = true;
                last_processed_id = result.frame_id;
            }
        }

        std::vector<dlib::rectangle> faces;
        std::vector<std::string> identities;
        if (new_result) 
        {
            faces = result.faces;
            identities = result.identities;
            last_faces = faces;
            last_identities = identities;
        } 
        else 
        {
            faces = last_faces;
            identities = last_identities;
        }

        bool should_recognize = (frame_count % sample_interval == 0);
        if (should_recognize && frame_count != last_processed_id) 
        {
            RecognitionTask task{frame_count, dlib_img};
            {
                std::lock_guard<std::mutex> lock(task_mutex);
                if (task_queue.size() < 5) 
                    task_queue.push(std::move(task));
            }
            task_cv.notify_one();
        }

        auto now = std::chrono::system_clock::now();
        std::lock_guard<std::mutex> vlock(visitor_mutex);
        std::unordered_set<std::string> current_frame_visitors;
        for (size_t i = 0; i < identities.size(); ++i) 
        {
            const auto& identity = identities[i];
            if (identity == "Stranger" || identity == "Error") 
                continue;
            current_frame_visitors.insert(identity);
            if (visitor_log.find(identity) == visitor_log.end()) 
            {
                VisitorInfo info;
                info.first_seen = now;
                info.last_seen = now;
                info.departure_time = now;
                info.present = true;

                std::time_t first_seen_t = std::chrono::system_clock::to_time_t(now);
                std::stringstream timestamp_ss;
                timestamp_ss << std::put_time(std::localtime(&first_seen_t), "%Y%m%d_%H%M%S");
                std::string timestamp = timestamp_ss.str();

                std::string img_path = "visitor_images/" + identity + "_" + timestamp + ".jpg";
                const auto& face = faces[i];
                cv::Rect face_rect(cv::Point(face.left(), face.top()), cv::Point(face.right(), face.bottom()));
                cv::Mat face_img = frame(face_rect & cv::Rect(0, 0, frame.cols, frame.rows));
                cv::imwrite(img_path, face_img);
                info.img_path = img_path;
                visitor_log[identity] = info;

                std::cout << "[访客登记] " << identity << " 首次访问于 "
                          << std::put_time(std::localtime(&first_seen_t), "%Y-%m-%d %H:%M:%S")
                          << ", 截图保存为: " << img_path << std::endl;
            } 
            else 
            {
                visitor_log[identity].last_seen = now;
                visitor_log[identity].present = true;
            }
        }

        for (auto& [name, info] : visitor_log) 
        {
            if (info.present && current_frame_visitors.find(name) == current_frame_visitors.end())
            {
                info.present = false;
                info.departure_time = now;
                std::time_t departure_t = std::chrono::system_clock::to_time_t(now);
                std::cout << "[访客离开] " << name << " 离开于 "
                          << std::put_time(std::localtime(&departure_t), "%Y-%m-%d %H:%M:%S")
                          << std::endl;
            }
        }

        for (size_t i = 0; i < faces.size(); ++i) 
        {
            const auto& face = faces[i];
            const auto& identity = identities[i];
            cv::Point top_left(face.left(), face.top());
            cv::Point bottom_right(face.right(), face.bottom());
            cv::Scalar color = (identity == "Stranger") ? cv::Scalar(0, 0, 255) :
                               (identity == "Error") ? cv::Scalar(255, 0, 0) :
                                                        cv::Scalar(0, 255, 0);
            cv::rectangle(frame, top_left, bottom_right, color, 2);
            cv::Point text_pos(face.left(), face.top() - 10);
            cv::putText(frame, identity, text_pos, cv::FONT_HERSHEY_SIMPLEX, 0.7, color, 2);
        }

        timer_manager.start("推流");
        std::vector<uchar> buff_bgr;
        cv::imencode(".jpg", frame, buff_bgr, params);
        streamer.publish("/bgr", std::string(buff_bgr.begin(), buff_bgr.end()));
        timer_manager.end("推流");
        timer_manager.end("帧捕获和处理全过程");

        double frame_time = timer_manager.getLastDuration("帧捕获和处理全过程");
        int delay = std::max(1, static_cast<int>(33.3 - frame_time));
        frame_count++;
        if (cv::waitKey(delay) >= 0) break;
    }

    stop_threads.store(true);
    task_cv.notify_all();
    for (auto& worker : workers) if (worker.joinable()) worker.join();
    if (cap.isOpened()) cap.release();
    streamer.stop();

    std::ofstream visitor_csv("visitor_log.csv", std::ios::out);
    visitor_csv << "Name,FirstSeen,LeaveTime,ImagePath\n";
    for (const auto& [name, info] : visitor_log) {
        std::time_t first = std::chrono::system_clock::to_time_t(info.first_seen);
        std::time_t depart = std::chrono::system_clock::to_time_t(info.departure_time);
        visitor_csv << name << ","
                    << std::put_time(std::localtime(&first), "%Y-%m-%d %H:%M:%S") << ","
                    << std::put_time(std::localtime(&depart), "%Y-%m-%d %H:%M:%S") << ","
                    << info.img_path << "\n";
    }
    visitor_csv.close();

    std::cout << "\n性能统计信息:" << std::endl;
    std::cout << "=============================================" << std::endl;
    timer_manager.printAllSummaries();
    return 0;
}